# coding:utf-8
import torch
import torch.nn as nn
from transformers import BertModel, AdamW
from ..config.config import Config
 
class CasRel(nn.Module):
    def __init__(self, conf):
        super().__init__()
        self.config = Config()
        self.bert = BertModel.from_pretrained(self.config.bert_path)
        # 定义第一个线性层，来判断主实体的头部位置
        self.sub_heads_linear = nn.Linear(self.config.bert_dim, 1)
        # 定义第二个线性层，来判断主实体的尾部位置
        self.sub_tails_linear = nn.Linear(self.config.bert_dim, 1)
        # 定义第三个线性层，来判断客实体的头部位置以及关系类型
        self.obj_heads_linear = nn.Linear(self.config.bert_dim, conf.num_rel)
        # 定义第四个线性层，来判断客实体的尾部位置以及关系类型
        self.obj_tails_linear = nn.Linear(self.config.bert_dim, conf.num_rel)
 
    def get_encoded_text(self, token_ids, mask):
        # 用于编码
        # token_ids 
        encoded_text = self.bert(token_ids, attention_mask=mask)[0]
        return encoded_text
 
    def get_subs(self, encoded_text):
        # 用来预测主实体开始索引和结束索引
        pre_sub_heads = torch.sigmoid(self.sub_heads_linear(encoded_text))
        pre_sub_tails = torch.sigmoid(self.sub_tails_linear(encoded_text))
        return pre_sub_heads, pre_sub_tails
 
    def get_objs_for_specific_sub(self, sub_head2tail, sub_len, encoded_text):
        '''
        将subject实体信息融合原始句子中：将主实体字向量实现平均，然后加在当前句子的每一个字向量上，进行计算
        :param sub_head2tail:shape-->【16，1, 200】
        :param sub_len:shape--->[16,1]
        :param encoded_text:.shape[16，200，768]
        :return:
            pred_obj_heads-->shape []
            pre_obj_tails-->shape []
        '''
        sub = torch.matmul(sub_head2tail, encoded_text)# 将主实体特征和编码后的文本进行融合
        sub_len = sub_len.unsqueeze(1) # 主实体长度（扩维）
        sub = sub / sub_len # 平均主实体信息
        encoded_text = encoded_text + sub #将处理后的实体特征和原始编码后的文本进行融合
        pred_obj_heads = torch.sigmoid(self.obj_heads_linear(encoded_text))
        pre_obj_tails = torch.sigmoid(self.obj_tails_linear(encoded_text))
        return pred_obj_heads, pre_obj_tails
 
    def forward(self, input_ids, mask, sub_head2tail, sub_len):
        '''
        :param input_ids: shape-->[16, 200]
        :param mask: shape-->[16, 200]
        :param sub_head2tail: shape-->[16, 200]
        :param sub_len: shape-->[16, 1]
        :return:
        '''
        # todo: encode_text.shape--->[16,200,768]
        encoded_text = self.get_encoded_text(input_ids, mask)
        pred_sub_heads, pre_sub_tails = self.get_subs(encoded_text)
        sub_head2tail = sub_head2tail.unsqueeze(1)
        pred_obj_heads, pre_obj_tails =self.get_objs_for_specific_sub(sub_head2tail, sub_len,encoded_text)
        result_dict = {'pred_sub_heads': pred_sub_heads,
                       'pred_sub_tails': pre_sub_tails,
                       'pred_obj_heads': pred_obj_heads,
                       'pred_obj_tails': pre_obj_tails,
                       'mask': mask}
        return result_dict
 
    def compute_loss(self,
                     pred_sub_heads, pred_sub_tails,
                     pred_obj_heads, pred_obj_tails,
                     mask,
                     sub_heads, sub_tails,
                     obj_heads, obj_tails):
        '''
        计算损失
        :param pred_sub_heads:[16, 200, 1]
        :param pred_sub_tails:[16, 200, 1]
        :param pred_obj_heads:[16, 200, 18]
        :param pred_obj_tails:[16, 200, 18]
        :param mask: shape-->[16, 200]
        :param sub_heads: shape-->[16, 200]
        :param sub_tails: shape-->[16, 200]
        :param obj_heads: shape-->[16, 200, 18]
        :param obj_tails: shape-->[16, 200, 18]
        :return:
        '''
        # todo:sub_heads.shape,sub_tails.shape, mask-->[16, 200]
        # todo:obj_heads.shape,obj_tails.shape-->[16, 200, 18]
        rel_count = obj_heads.shape[-1]
        rel_mask = mask.unsqueeze(-1).repeat(1, 1, rel_count)
        loss_1 = self.loss(pred_sub_heads, sub_heads, mask)
        loss_2 = self.loss(pred_sub_tails, sub_tails, mask)
        loss_3 = self.loss(pred_obj_heads, obj_heads, rel_mask)
        loss_4 = self.loss(pred_obj_tails, obj_tails, rel_mask)
        return loss_1 + loss_2 + loss_3 + loss_4
 
    def loss(self, pred, gold, mask):
        pred = pred.squeeze(-1)
        los = nn.BCELoss(reduction='none')(pred, gold)
        if los.shape != mask.shape:
            mask = mask.unsqueeze(-1)
        los = torch.sum(los * mask) / torch.sum(mask)
        return los
 
 
